{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMSoUMKic3BwTFLpQ1YCwEm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sugarforever/LangChain-Tutorials/blob/main/Langchain_Memory_Persistent_Store.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 简介\n",
        "\n",
        "该Python Notebook演示LangChain框架所提供的Memory持久化能力。LangChain通过langchain.memory.chat_message_histories包中提供的一系列组件支持多种形式的历史消息存储，包括文件，数据库等。\n",
        "\n",
        "本示例演示文件存储形式的持久化。"
      ],
      "metadata": {
        "id": "q0ApeCaVrZhV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install langchain openai --quiet --upgrade"
      ],
      "metadata": {
        "id": "RdSvP1Oi7LCc"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "CcSMVHBz6zrj"
      },
      "outputs": [],
      "source": [
        "from langchain.llms import OpenAI\n",
        "from langchain.chains import ConversationChain\n",
        "from langchain.callbacks import get_openai_callback\n",
        "from langchain.memory import ConversationBufferMemory"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. 准备**OpenAI API Key**，**track_tokens_usage**辅助函数（调用`chain.run`函数，并统计token开销），以及LLM实例"
      ],
      "metadata": {
        "id": "wVqmb_UYrif9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "os.environ['OPENAI_API_KEY'] = 'your openai api key'"
      ],
      "metadata": {
        "id": "E2xdWZRTiKWV"
      },
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def track_tokens_usage(chain, query):\n",
        "    with get_openai_callback() as cb:\n",
        "        result = chain.run(query)\n",
        "        print(f'Total tokens: {cb.total_tokens}')\n",
        "\n",
        "    return result\n"
      ],
      "metadata": {
        "id": "yozbiUy2_RCt"
      },
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "llm = OpenAI(\n",
        "  temperature=0,\n",
        "\topenai_api_key=os.environ[\"OPENAI_API_KEY\"],\n",
        "\tmodel_name=\"text-davinci-003\"\n",
        ")\n"
      ],
      "metadata": {
        "id": "akq_I-dn-hpd"
      },
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "2. 使用Memory组件**ConversationBufferMemory**的默认配置"
      ],
      "metadata": {
        "id": "EU9QeEcZr-x1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "conversation = ConversationChain(llm=llm, memory = ConversationBufferMemory())"
      ],
      "metadata": {
        "id": "czgOmwfsr9Y7"
      },
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(conversation.prompt.template)"
      ],
      "metadata": {
        "id": "pEsc6-AF-riz",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0de9b559-782e-49ea-bf4f-b1e437eb7f1e"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
            "\n",
            "Current conversation:\n",
            "{history}\n",
            "Human: {input}\n",
            "AI:\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "track_tokens_usage(conversation, \"What is Langchain?\")\n"
      ],
      "metadata": {
        "id": "rM_qZkMT_nw7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 73
        },
        "outputId": "cfc70405-94b9-4fb1-9229-0849811e01e1"
      },
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Total tokens: 110\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "' Langchain is a blockchain-based language learning platform that allows users to learn new languages in a secure and decentralized environment. It uses a combination of blockchain technology and artificial intelligence to provide users with an immersive language learning experience.'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "3. 引入消息历史持久化组件"
      ],
      "metadata": {
        "id": "OnAq2L8rsLLy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.memory.chat_message_histories import FileChatMessageHistory"
      ],
      "metadata": {
        "id": "TymjTQMWiXAx"
      },
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "message_history = FileChatMessageHistory(file_path = 'conversation_20230620.txt')\n",
        "\n",
        "memory = ConversationBufferMemory(chat_memory=message_history)"
      ],
      "metadata": {
        "id": "EL-gMHvdjLHD"
      },
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "conversation = ConversationChain(llm=llm, memory = memory)"
      ],
      "metadata": {
        "id": "9K799NAZjVUT"
      },
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "track_tokens_usage(conversation, \"How does ChatGPT enable coherent conversation?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 73
        },
        "id": "D1kDz1EjjYcr",
        "outputId": "bdd9c106-7999-48f8-e4f1-815c13925bc7"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Total tokens: 133\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "' ChatGPT is a natural language processing (NLP) model that enables coherent conversation by understanding the context of a conversation and providing relevant responses. It uses a transformer-based architecture to generate responses that are based on the context of the conversation. It also uses a large corpus of data to generate more accurate responses.'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cat conversation_20230620.txt"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U8rrP-T3kOS0",
        "outputId": "c03707e3-0a90-4e2b-bd47-571372d1f148"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[{\"type\": \"human\", \"data\": {\"content\": \"How does ChatGPT enable coherent conversation?\", \"additional_kwargs\": {}, \"example\": false}}, {\"type\": \"ai\", \"data\": {\"content\": \" ChatGPT is a natural language processing (NLP) model that enables coherent conversation by understanding the context of a conversation and providing relevant responses. It uses a transformer-based architecture to generate responses that are based on the context of the conversation. It also uses a large corpus of data to generate more accurate responses.\", \"additional_kwargs\": {}, \"example\": false}}]"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "track_tokens_usage(conversation, \"Bye now!\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        },
        "id": "oPBsl3jUyNhP",
        "outputId": "a2e54fdc-c9ca-4c2f-e67f-9c159552bf36"
      },
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Total tokens: 156\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "' Bye! It was nice talking to you. Have a great day!'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cat conversation_20230620.txt"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mjGRBf5DyO7I",
        "outputId": "557083c1-4c41-4515-b685-90e0b269d153"
      },
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[{\"type\": \"human\", \"data\": {\"content\": \"How does ChatGPT enable coherent conversation?\", \"additional_kwargs\": {}, \"example\": false}}, {\"type\": \"ai\", \"data\": {\"content\": \" ChatGPT is a natural language processing (NLP) model that enables coherent conversation by understanding the context of a conversation and providing relevant responses. It uses a transformer-based architecture to generate responses that are based on the context of the conversation. It also uses a large corpus of data to generate more accurate responses.\", \"additional_kwargs\": {}, \"example\": false}}, {\"type\": \"human\", \"data\": {\"content\": \"Bye now!\", \"additional_kwargs\": {}, \"example\": false}}, {\"type\": \"ai\", \"data\": {\"content\": \" Bye! It was nice talking to you. Have a great day!\", \"additional_kwargs\": {}, \"example\": false}}]"
          ]
        }
      ]
    }
  ]
}